{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `GPUDrive` simulator concepts\n",
    "\n",
    "In this notebook, we demonstrate how to work with the `GPUDrive` simulator and access its basic attributes in Python. The simulator, written in C++, is built on top of the [Madrona Engine](https://madrona-engine.github.io/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from pathlib import Path\n",
    "import madrona_gpudrive\n",
    "\n",
    "# Set working directory to the base directory 'gpudrive_madrona'\n",
    "working_dir = Path.cwd()\n",
    "while working_dir.name != 'gpudrive':\n",
    "    working_dir = working_dir.parent\n",
    "    if working_dir == Path.home():\n",
    "        raise FileNotFoundError(\"Base directory 'gpudrive_madrona' not found\")\n",
    "os.chdir(working_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Summary\n",
    "\n",
    "- `GPUDrive` simulations are discretized traffic scenarios. A scenario is a constructed snapshot of traffic situation at a particular timepoint.\n",
    "- The state of the vehicle of focus is referred to as the **ego state**. Each vehicle has their own partial view of the traffic scene; and a visible state is constructed by parameterizing the view distance of the driver. \n",
    "- The `step()` method advances the simulation with a desired step size. By default, the dynamics of vehicles are driven by a kinematic bicycle model. If a vehicle is not controlled (that is, we do not give it actions), its position, heading, and speed will be updated according to a the log replay demonstrations.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiating a sim object with default parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = madrona_gpudrive.Parameters()\n",
    "params.polylineReductionThreshold = 0.1\n",
    "\n",
    "sim = madrona_gpudrive.SimManager(\n",
    "    exec_mode=madrona_gpudrive.madrona.ExecMode.CUDA\n",
    "    if device == \"cuda\"\n",
    "    else madrona_gpudrive.madrona.ExecMode.CPU,\n",
    "    gpu_id=0,\n",
    "    scenes=[\"data/processed/examples/tfrecord-00000-of-01000_325.json\"],\n",
    "    params=params,  # Environment parameters\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The simulator provides the following functions:\n",
    "- `reset(world_idx)` resets a specific world or environment at the given index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.reset([0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- `step()` advances the dynamics of all worlds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exporting tensors\n",
    "\n",
    "To retrieve a tensor from the simulator, call the specific `tensor()` method, followed by either `to_torch()` or `to_jax()`.\n",
    "\n",
    "For example, here is how to access the ego state, or self-observation tensor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 64, 8]), device(type='cpu'))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "observation_tensor = sim.self_observation_tensor().to_torch()\n",
    "\n",
    "observation_tensor.shape, observation_tensor.device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or alternatively:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "observation_tensor_jax = sim.self_observation_tensor().to_jax()\n",
    "\n",
    "observation_tensor_jax.shape, observation_tensor_jax.devices()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are all available tensor exports and methods on the sim object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "absolute_self_observation_tensor\n",
      "action_tensor\n",
      "agent_roadmap_tensor\n",
      "controlled_state_tensor\n",
      "deleteAgents\n",
      "deleted_agents_tensor\n",
      "depth_tensor\n",
      "done_tensor\n",
      "expert_trajectory_tensor\n",
      "info_tensor\n",
      "lidar_tensor\n",
      "map_name_tensor\n",
      "map_observation_tensor\n",
      "metadata_tensor\n",
      "partner_observations_tensor\n",
      "reset\n",
      "response_type_tensor\n",
      "reward_tensor\n",
      "rgb_tensor\n",
      "self_observation_tensor\n",
      "set_maps\n",
      "shape_tensor\n",
      "step\n",
      "steps_remaining_tensor\n",
      "valid_state_tensor\n",
      "world_means_tensor\n"
     ]
    }
   ],
   "source": [
    "for attr in dir(sim):\n",
    "    if not attr.startswith(\"_\"):\n",
    "        print(attr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inspect valid and controlled agents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To check the number of agents and road points in each world, you can use the `shape_tensor`.\n",
    "\n",
    "The shape tensor is a 2D tensor where the first dimension represents the number of worlds, and the second dimension represents the shape of each world."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shape_tensor = sim.shape_tensor().to_jax()\n",
    "print(f\"Shape tensor has a shape of (Num Worlds, 2): {shape_tensor.shape}\")\n",
    "\n",
    "for world_idx in range(shape_tensor.shape[0]):\n",
    "    print(\n",
    "        f\"World {world_idx} has {shape_tensor[world_idx][0]} VALID agents and {shape_tensor[world_idx][1]} VALID road objects\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The number of **valid** agents indicates the number of controllable agents (vehicles). Some vehicles or bicycles may be initialized in incorrect positions or remain static; these are marked as **invalid** and cannot be controlled.\n",
    "\n",
    "The sim comes with a mask that indicates which agents can be controlled. Entries are `1` for agents that can be controlled, and `0` otherwise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Controlled state tensor has a shape of (num_worlds, max_num_agents_in_scene, 1):  torch.Size([1, 64, 1])\n"
     ]
    }
   ],
   "source": [
    "controlled_state_tensor = sim.controlled_state_tensor().to_torch()\n",
    "print(\n",
    "    \"Controlled state tensor has a shape of (num_worlds, max_num_agents_in_scene, 1): \",\n",
    "    controlled_state_tensor.shape,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype=torch.int32)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# We can control 3 agents in this world\n",
    "controlled_state_tensor.squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "controlled_state_tensor.sum().item()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Actions\n",
    "\n",
    "The action space is determined by the chosen dynamics_model. By default, it uses a bicycle model with the following tuple of actions:\n",
    "\n",
    "- **Acceleration**: Continuous float values representing the acceleration applied to the agents. This affects how quickly an agent speeds up or slows down.\n",
    "- **Steering Angle**: Continuous float values representing the steering angle, following the bicycle kinematic model. This determines how sharply an agent turns.\n",
    "- **Heading Angle**: Continuous float values for the heading angle, which control the direction an agent is facing.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The action tensor stores the current actions for all agents across all worlds:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Action tensor has a shape of (num_worlds, max_num_agents_in_scene, 3): torch.Size([1, 64, 10])\n"
     ]
    }
   ],
   "source": [
    "action_tensor = sim.action_tensor().to_torch()\n",
    "print(\n",
    "    f\"Action tensor has a shape of (num_worlds, max_num_agents_in_scene, 3): {action_tensor.shape}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To set the actions for all controlled agents, we use the `copy_()` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Actions tensor after setting all actions to 1: tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])\n"
     ]
    }
   ],
   "source": [
    "actions_tensor = sim.action_tensor().to_torch()\n",
    "\n",
    "actions = torch.full(actions_tensor.shape, 1.0)\n",
    "actions_tensor.copy_(actions)\n",
    "\n",
    "print(f\"Actions tensor after setting all actions to 1: {actions_tensor[0][0]}\")\n",
    "\n",
    "# Call step() to apply the actions\n",
    "sim.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inspecting the simulator settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameters:\n",
      "IgnoreNonVehicles   : False\n",
      "collisionBehaviour  : madrona_gpudrive.CollisionBehaviour.AgentStop\n",
      "disableClassicalObs : False\n",
      "dynamicsModel       : madrona_gpudrive.DynamicsModel.Classic\n",
      "enableLidar         : False\n",
      "initOnlyValidAgentsAtFirstStep: True\n",
      "isStaticAgentControlled: False\n",
      "maxNumControlledAgents: 10000\n",
      "observationRadius   : 0.0\n",
      "polylineReductionThreshold: 0.0\n",
      "rewardParams        : <madrona_gpudrive.RewardParams object at 0x7f9d1c273ab0>\n",
      "Reward parameters:\n",
      "    distanceToExpertThreshold: 0.0\n",
      "    distanceToGoalThreshold: 0.0\n",
      "    rewardType        : madrona_gpudrive.RewardType.DistanceBased\n",
      "roadObservationAlgorithm: madrona_gpudrive.FindRoadObservationsWith.KNearestEntitiesWithRadiusFiltering\n"
     ]
    }
   ],
   "source": [
    "params = madrona_gpudrive.Parameters()\n",
    "\n",
    "print(\"Parameters:\")\n",
    "for attr in dir(params):\n",
    "    if not attr.startswith(\"__\"):\n",
    "        value = getattr(params, attr)\n",
    "        print(f\"{attr:20}: {value}\")\n",
    "        if attr == \"rewardParams\":\n",
    "            print(\"Reward parameters:\")\n",
    "            reward_params = getattr(params, attr)\n",
    "            for attr2 in dir(reward_params):\n",
    "                if not attr2.startswith(\"__\"):\n",
    "                    value2 = getattr(reward_params, attr2)\n",
    "                    print(f\"    {attr2:18}: {value2}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuring the simulator "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To set the parameters of the simulator, fill in the values for each attribute of the parameter object as below. This allows you to customize the simulation settings.\n",
    "\n",
    "The params object can be passed to the sim constructor like this:\n",
    "\n",
    "```Python\n",
    "sim = gpudrive_madrona.SimManager(\n",
    "    ...\n",
    "    params=params \n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "reward_params = madrona_gpudrive.RewardParams()\n",
    "reward_params.rewardType = madrona_gpudrive.RewardType.DistanceBased\n",
    "reward_params.distanceToGoalThreshold = 1.0\n",
    "reward_params.distanceToExpertThreshold = 1.0\n",
    "\n",
    "# Initialize Parameters\n",
    "params = madrona_gpudrive.Parameters()\n",
    "params.polylineReductionThreshold = 1.0\n",
    "params.observationRadius = 100.0\n",
    "params.collisionBehaviour = madrona_gpudrive.CollisionBehaviour.Ignore\n",
    "params.maxNumControlledAgents = 10\n",
    "params.rewardParams = reward_params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step simulator and access information\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "actions_shape = sim.action_tensor().to_torch().shape\n",
    "dones = sim.done_tensor().to_torch()\n",
    "\n",
    "\n",
    "obs, rews, dones = (\n",
    "    sim.self_observation_tensor().to_torch(),\n",
    "    sim.reward_tensor().to_torch(),\n",
    "    sim.done_tensor().to_torch(),\n",
    ")\n",
    "actions = torch.rand(actions_shape)\n",
    "sim.action_tensor().to_torch().copy_(actions)\n",
    "sim.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here’s your text with all the headings formatted as lowercase (except for the first letter):\n",
    "\n",
    "---\n",
    "\n",
    "## Simulator configurations and objects\n",
    "\n",
    "This outlines the key components of the simulator, including observation tensors, rewards, road reduction algorithms, collision behavior, and miscellaneous parameters.\n",
    "\n",
    "### Observation space\n",
    "\n",
    "- You can use the datatypes in [`gpudrive/datatypes/observations.py`](https://github.com/Emerge-Lab/gpudrive/blob/main/pygpudrive/datatypes/observation.py) to easily index the tensors below. More documentation is provided there.\n",
    "\n",
    "#### **SelfObservation**\n",
    "\n",
    "The `SelfObservation` tensor of shape `(6,)` for each agent represents the agent's state. The respective values are:\n",
    "\n",
    "- `SelfObservation[0]`: Current *speed* of the agent.\n",
    "- `SelfObservation[1]`: *Length* of the agent.\n",
    "- `SelfObservation[2]`: *Width* of the agent.\n",
    "- `SelfObservation[3:4]`: *Coordinates (x, y)* of the goal relative to the agent.\n",
    "- `SelfObservation[5]`: Collision state of the agent (`0` for no collision, `1` for collision).\n",
    "\n",
    "#### **MapObservation**\n",
    "\n",
    "The `MapObservation` tensor of shape `(8,)` provides absolute position data for map objects. Values include:\n",
    "\n",
    "- `MapObservation[0:1]`: Position `(x, y)` of the map object.\n",
    "- `MapObservation[2:4]`: Scale `(length, width, height)` of the map object.\n",
    "- `MapObservation[5]`: Heading angle of the map object.\n",
    "- `MapObservation[6]`: Type of the map object.\n",
    "- `MapObservation[7]`: Road ID (`>0` for valid ID, `-1` for unassigned, `0` for missing data).\n",
    "\n",
    "#### **PartnerObservation**\n",
    "\n",
    "The `PartnerObservation` tensor of shape `(num_agents-1, 7)` provides data about other agents relative to the ego agent within the `observationRadius`:\n",
    "\n",
    "- `PartnerObservation[0]`: *Speed* of the neighboring agent.\n",
    "- `PartnerObservation[1:2]`: *Relative position (x, y)* of the neighboring agent.\n",
    "- `PartnerObservation[3]`: *Orientation* of the neighboring agent.\n",
    "- `PartnerObservation[4:5]`: *Dimensions (length, width)* of the neighboring agent.\n",
    "- `PartnerObservation[6]`: Type of the neighboring agent.\n",
    "\n",
    "#### **AgentMapObservations**\n",
    "\n",
    "The `AgentMapObservations` tensor of shape `(num_road_objs, 8)` provides road object data relative to the ego agent. The values are identical to `MapObservation`.\n",
    "\n",
    "### Rewards\n",
    "\n",
    "- **Reward types**:\n",
    "  - `DistanceBased`: Distance from the agent to the goal.\n",
    "  - `OnGoalAchieved`: Binary reward (`1` for reaching goal, `0` otherwise) (recommended).\n",
    "  - `Dense`: Distance from the expert trajectory.\n",
    "- **Thresholds**:\n",
    "  - `distanceToGoalThreshold`: Default: `0.0`.\n",
    "  - `distanceToExpertThreshold`: Default: `0.0`.\n",
    "\n",
    "### Road graph \n",
    "\n",
    "- You can use the datatypes in [`gpudrive/datatypes/roadgraph.py`](https://github.com/Emerge-Lab/gpudrive/blob/main/gpudrive/datatypes/roadgraph.py) to easily index the tensors below. More documentation is provided there.\n",
    "\n",
    "#### Road reduction algorithm\n",
    "\n",
    "The **Visvalingam-Whyatt Algorithm** is applied to simplify road lines and edges.\n",
    "\n",
    "- **Parameter**:  \n",
    "  `polylineReductionThreshold`: Ranges from `0` (no reduction) to `1`. Default: `0.5`.\n",
    "\n",
    "### Collision behavior\n",
    "\n",
    "Three options for the collision behavior are available. These determine what happens when an agent touches a road edge or other agents' bounding box.\n",
    "\n",
    "- `AgentStop`: Agents stop upon collision.\n",
    "- `AgentRemoved`: Agents are removed from the scene upon collision.\n",
    "- `Ignore`: Collision is logged but does not affect the agent's movement.\n",
    "\n",
    "### Miscellaneous parameters\n",
    "\n",
    "- `ObservationRadius`: Radius around every agent that determines visible view. Elements outside this radius are zeroed out.\n",
    "- `MaxNumControlledVehicles`: Maximum number of controlled agents. Default: First valid agents are controlled.\n",
    "- `IgnoreNonVehicles`: Excludes non-vehicle entities (e.g., pedestrians, cyclists). Default: `false`.\n",
    "- `roadObservationAlgorithm`: Options:\n",
    "  - `KNearestEntitiesWithRadiusFiltering`: Limits observations to the nearest points within `observationRadius`.\n",
    "  - `AllEntitiesWithRadiusFiltering`: Searches all points within `observationRadius`. Default: `KNearestEntitiesWithRadiusFiltering`.\n",
    "- `initOnlyValidAgentsAtFirstStep`: Initializes only valid agents. Default: `true`.\n",
    "- `isStaticAgentControlled`: Controls if parked agents can be set static. Default: `false`.\n",
    "- `enableLidar`: Enables lidar observations.\n",
    "- `disableClassicalObs`: Disables `PartnerObservations` and `AgentMapObservations` for performance. Default: `false`.\n",
    "\n",
    "### Object types\n",
    "\n",
    "**Road types**:\n",
    "- `RoadEdge`, `RoadLine`, `RoadLane`, `CrossWalk`, `SpeedBump`, `StopSign`.\n",
    "\n",
    "**Agent types**:\n",
    "- `Vehicle`, `Pedestrian`, `Cyclist`.\n",
    "\n",
    "**Other types**:\n",
    "- `Padding`: Ensures consistent output shapes.\n",
    "- `None`: Marks entities as invalid (e.g., out of `ObservationRadius` or removed due to collision).\n",
    "\n",
    "For any further details or indexing references, please explore [`gpudrive/datatypes`](https://github.com/Emerge-Lab/gpudrive/tree/main/gpudrive/datatypes) module."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Other exported tensors\n",
    "\n",
    "- By default, the x, y, z coordinates are de-meaned to center everything at 0. To revert back to the original coordinates, we export the `world_means_tensor`. \n",
    "    - Shape:  (`num_worlds`, 3)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
